#coding=utf-8
import  cv2
import numpy as np
import matplotlib.pyplot as plt

trainData = np.random.randint(0,100,(50,2)).astype(np.float32)
responses = np.random.randint(0,2,(50,1)).astype(np.float32)

#print responses.ravel() == 0
red = trainData[responses.ravel() == 0]#将数组中为0的数据标记为红色
plt.scatter(red[:,0],red[:,1],80,'r','^')#xy 大小 颜色  形状

blue = trainData[responses.ravel() == 1]
plt.scatter(blue[:,0],blue[:,1],80,'b','s')

newcomer = np.random.randint(0,100,(10,2)).astype(np.float32)
plt.scatter(newcomer[:,0],newcomer[:,1],80,'g','o')

knn = cv2.KNearest()
knn.train(trainData,responses)
ret,result,neighbours,dist = knn.find_nearest(newcomer,3)

print 'result = ',result
print 'neighbours = ',neighbours
print 'distance = ', dist

plt.show()